import json
import math
import triton_python_backend_utils as pb_utils
import cv2
import onnxruntime
import numpy as np
import os
import logging
from logging.handlers import RotatingFileHandler
import onnxruntime as ort


log_filename = '/model.log'
ort.set_default_logger_severity(3)
# Create handlers
file_handler = RotatingFileHandler(log_filename, maxBytes=50 * 1024 * 1024, backupCount=5)
console_handler = logging.StreamHandler()

# Create formatters and add them to handlers
log_formatter = logging.Formatter('%(asctime)s - %(levelname)s - %(message)s')
file_handler.setFormatter(log_formatter)
console_handler.setFormatter(log_formatter)

# Create logger, add handlers
logger = logging.getLogger()
logger.setLevel(logging.INFO)
logger.addHandler(file_handler)
logger.addHandler(console_handler)


# Function to add a separator in the logs
def add_log_separator():
    logger.info("-" * 80)


local_path = os.path.dirname(os.path.abspath(__file__))
model_filename = "model.onnx"
onnx_model_path = os.path.join(local_path, model_filename)


class TritonPythonModel:
    def initialize(self, args):
        try:
            if args is None:
                raise ValueError("Received 'None' for 'args'")
            if 'model_config' not in args:
                raise ValueError("Expected 'model_config' in 'args', but not found")

            self.model_config = json.loads(args['model_config'])

            out_bboxes_config = pb_utils.get_output_config_by_name(self.model_config, "bboxes")
            out_scores_config = pb_utils.get_output_config_by_name(self.model_config, "scores")

            self.out_bboxes_dtype = pb_utils.triton_string_to_numpy(out_bboxes_config['data_type'])
            self.out_scores_dtype = pb_utils.triton_string_to_numpy(out_scores_config['data_type'])
            logger.info(f"The model is loading...")
            self.sess = onnxruntime.InferenceSession(onnx_model_path)
            logging.info("Model loaded successfully!")
        except Exception as e:
            logging.error(f'Failed to initialize model: {e}')
            raise

    def execute(self, requests):
        responses = []
        for request in requests:
            try:
                add_log_separator()
                logging.info(f"Received request...")
                image = pb_utils.get_input_tensor_by_name(request, 'image').as_numpy()  # (337, 600, 3)
                score = pb_utils.get_input_tensor_by_name(request, 'score')
                if score is None:
                    score = np.float32([0.3])
                    logger.info("No score input, use default value 0.3")
                else:
                    score = score
                    score = score.as_numpy()
                    score = score.astype(np.float32)
                logging.info(f"Input image shape: {image.shape}, score: {score}")
                # [[188 126 250 214][39 187 92 232]]  [[0.91708493 0.88484555]]    [[1 1]]
                result, img_after = self.Inference(image, score)

                # whether the result is empty
                if result[0][0].size == 0:
                    logger.info("No object detected!!!")
                    e_bboxes = np.empty((0, 4), dtype=self.out_bboxes_dtype)
                    e_conf = np.empty((0, ), dtype=self.out_scores_dtype)
                   
                    out_tensor_bboxes = pb_utils.Tensor('bboxes', e_bboxes)
                    out_tensor_scores = pb_utils.Tensor('scores', e_conf)
                    
                    inference_response = pb_utils.InferenceResponse(
                        output_tensors=[out_tensor_bboxes, out_tensor_scores])
                    responses.append(inference_response)
                    continue
                bboxes = self.cod_trf(result[0], image, img_after)
                conf = np.array(result[1][0] * 100)
                bbox = np.array(np.round(bboxes).astype(np.uint32))
                logger.info(f"The number of detected targets: {len(bbox)}")
                logger.info(f"Inference results: bboxes:{bbox} conf:{conf}")

                # 创建输出张量对象
                out_tensor_bboxes = pb_utils.Tensor('bboxes', bbox.astype(self.out_bboxes_dtype))
                out_tensor_scores = pb_utils.Tensor('scores', conf.astype(self.out_scores_dtype))
                # 创建推理响应对象，包含输出张量
                inference_response = pb_utils.InferenceResponse(
                    output_tensors=[out_tensor_bboxes, out_tensor_scores])
                responses.append(inference_response)
            except Exception as e:
                logging.error(f'Failed to execute request: {e}')
                inference_response = pb_utils.InferenceResponse(output_tensors=[], error=pb_utils.TritonError(str(e)))
                responses.append(inference_response)
        return responses

    def finalize(self):
        logger.info('Cleaning up...')

    def Inference(self, image, score):
        try:
            std_h, std_w = 640, 640
            img_after = self.resize_image(image, (std_w, std_h), True)  # （640， 640， 3）
            imageData = self.img2input(img_after)

            
            result = self.sess.run(None, {'images': imageData, 'score': score})
            return result, img_after
        
        except Exception as e:
            logger.error(f"Inference failed: {e}")
            raise


    def resize_image(self, image, size, letterbox_image):
        ih, iw, _ = image.shape
        h, w = size
        if letterbox_image:
            scale = min(w / iw, h / ih)
            nw = int(iw * scale)
            nh = int(ih * scale)
            image = cv2.resize(image, (nw, nh), interpolation=cv2.INTER_LINEAR)
            image_back = np.ones((h, w, 3), dtype=np.uint8) * 128
            image_back[(h - nh) // 2: (h - nh) // 2 + nh, (w - nw) // 2:(w - nw) // 2 + nw, :] = image
        else:
            image_back = image
        return image_back

    def img2input(self, img):
        img = np.transpose(img, (2, 0, 1))
        img = img / 255
        return np.expand_dims(img, axis=0).astype(np.float32)

    def xywh2xyxy(self, *box):
        ret = [box[0] - box[2] // 2, box[1] - box[3] // 2, box[0] + box[2] // 2, box[1] + box[3] // 2]
        return ret

    def cod_trf(self, result, pre, after):
        result_boxes = []
        # result[0] = [[1,2,3,4], [2,3,6,5]]
        for i in result[0]:
            x, y, w, h = i
            x1, y1, x2, y2 = self.xywh2xyxy(x, y, w, h)
            h_pre, w_pre, _ = pre.shape
            h_after, w_after, _ = after.shape
            scale = max(w_pre / w_after, h_pre / h_after)
            h_pre, w_pre = h_pre / scale, w_pre / scale
            x_move, y_move = abs(w_pre - w_after) // 2, abs(h_pre - h_after) // 2
            ret_x1, ret_x2 = (x1 - x_move) * scale, (x2 - x_move) * scale
            ret_y1, ret_y2 = (y1 - y_move) * scale, (y2 - y_move) * scale
            ret = np.array([ret_x1, ret_y1, ret_x2, ret_y2])
            result_boxes.append(ret)
        result[0] = result_boxes
        return result_boxes
